const fs = require('fs');
const path = require('path');
const axios = require('axios');
const mysql = require('mysql2/promise');
require('dotenv').config({ path: './nestjs-app/.env' });

const OUTPUT_DIR = path.resolve(__dirname, 'arming');
const BASE_URL = 'http://cd521.xin/arming';

function ensureOutputDirectory(directoryPath) {
    if (!fs.existsSync(directoryPath)) {
        fs.mkdirSync(directoryPath, { recursive: true });
    }
}

async function downloadWithRetry(url, filePath, maxRetries) {
    let attempt = 0;
    let lastError = null;

    while (attempt <= maxRetries) {
        try {
            const response = await axios({ method: 'GET', url, responseType: 'stream', timeout: 15000 });
            await new Promise((resolve, reject) => {
                const writer = fs.createWriteStream(filePath);
                response.data.pipe(writer);
                writer.on('finish', resolve);
                writer.on('error', reject);
            });
            return;
        } catch (error) {
            lastError = error;
            attempt++;
            if (attempt > maxRetries) break;
            const delayMs = attempt * 1000;
            await new Promise(resolve => setTimeout(resolve, delayMs));
        }
    }

    throw new Error(`下载失败（重试${maxRetries}次后）：${lastError && lastError.message ? lastError.message : '未知错误'}`);
}

function createTaskRunner(maxConcurrency) {
    let running = 0;
    const queue = [];

    function runNext() {
        if (running >= maxConcurrency || queue.length === 0) return;
        const { task, resolve, reject } = queue.shift();
        running++;
        task().then(result => {
            running--;
            resolve(result);
            runNext();
        }).catch(err => {
            running--;
            reject(err);
            runNext();
        });
    }

    return function enqueue(task) {
        return new Promise((resolve, reject) => {
            queue.push({ task, resolve, reject });
            runNext();
        });
    };
}

async function fetchUniqueImages(connection) {
    const [rows] = await connection.execute(
        `SELECT DISTINCT image FROM equipment_base WHERE image IS NOT NULL AND image <> ''`
    );
    const images = rows
        .map(r => r.image)
        .filter(Boolean)
        .map(v => String(v).trim())
        .filter(v => v.length > 0);
    return Array.from(new Set(images));
}

async function main() {
    ensureOutputDirectory(OUTPUT_DIR);

    const connection = await mysql.createConnection({
        host: process.env.DB_HOST || '43.143.253.188',
        port: Number(process.env.DB_PORT || 3306),
        user: process.env.DB_USER || 'nextjs',
        password: process.env.DB_PASS || 'yEAiRE2jcxTyRBy5',
        database: process.env.DB_NAME || 'nextjs',
        timezone: '+08:00'
    });

    try {
        console.log('读取 equipment_base.image 列表...');
        const images = await fetchUniqueImages(connection);
        console.log(`共找到 ${images.length} 个唯一 image 记录。`);

        const enqueue = createTaskRunner(8);
        let successCount = 0;
        let skipCount = 0;
        let failCount = 0;
        const failures = [];

        const tasks = images.map((imageName, index) => async () => {
            const url = `${BASE_URL}/${encodeURIComponent(imageName)}.png`;
            const filePath = path.join(OUTPUT_DIR, `${imageName}.png`);

            if (fs.existsSync(filePath)) {
                skipCount++;
                if ((index + 1) % 50 === 0) console.log(`[进度 ${index + 1}/${images.length}] 已存在，跳过: ${imageName}.png`);
                return;
            }

            try {
                await downloadWithRetry(url, filePath, 2);
                successCount++;
                if ((index + 1) % 25 === 0) console.log(`[进度 ${index + 1}/${images.length}] 成功: ${imageName}.png`);
            } catch (err) {
                failCount++;
                failures.push({ image: imageName, url, error: err.message });
                console.log(`[进度 ${index + 1}/${images.length}] 失败: ${imageName}.png -> ${err.message}`);
            }
        });

        await Promise.all(tasks.map(task => enqueue(task)));

        console.log('\n=== 下载完成 ===');
        console.log(`成功: ${successCount}`);
        console.log(`跳过(已存在): ${skipCount}`);
        console.log(`失败: ${failCount}`);

        if (failures.length > 0) {
            const reportPath = path.join(OUTPUT_DIR, 'arming_download_failures.json');
            fs.writeFileSync(reportPath, JSON.stringify({ timestamp: new Date().toISOString(), failures }, null, 2), 'utf8');
            console.log(`失败详情已保存: ${reportPath}`);
        }
    } finally {
        await connection.end();
    }
}

main().catch(err => {
    console.error('执行出错:', err.message);
    process.exitCode = 1;
});


